import tensorflow as tf  
from tensorflow.keras.datasets import mnist  
from tensorflow.keras.layers import Flatten, Dense  
from tensorflow.keras.models import Sequential 
from tensorflow.keras.utils import to_categorical 
import numpy as np 
from tabulate import tabulate  
import matplotlib.pyplot as plt  


(X_train, y_train), (X_test, y_test) = mnist.load_data()  
X_train = X_train / 255.0  
X_test = X_test / 255.0  
y_train = to_categorical(y_train) 
y_test = to_categorical(y_test)  

# 2. 构建网络模型
model = Sequential()  # 创建Sequential模型
model.add(Flatten(input_shape=(28, 28)))  # 添加Flatten层,将二维图像数据展平为一维向量
model.add(Dense(20, activation=tf.nn.relu))  # 添加全连接层,使用ReLU激活函数,该层有20个神经元
model.add(Dense(10, activation=tf.nn.softmax))  # 添加输出层,使用Softmax激活函数,该层有10个神经元,对应10个类别

# 3. 编译模型,指定优化器、损失函数和评估指标
model.compile(optimizer='adam', loss='categorical_crossentropy', metrics=['accuracy'])

# 4. 训练模型
model.fit(X_train, y_train, epochs=20, validation_split=0.2)  # 训练模型,设置训练轮数为20,验证集比例为20%

# 5. 评估模型
test_loss, test_acc = model.evaluate(X_test, y_test)  # 在测试集上评估模型性能
print('Test accuracy:', test_acc)  # 输出测试集准确率

# 保存模型到硬盘
params_path = 'model_params.pth'
torch.save(model.state_dict(), params_path))